# 导包
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_curve, roc_auc_score
# 标准化
from sklearn.preprocessing import StandardScaler
# XGBoost
from xgboost import XGBClassifier
# 交叉网格
from sklearn.model_selection import GridSearchCV

# 读取数据
data = pd.read_csv('../data/train.csv')

# 删除无用列
data = data.drop(['StandardHours'], axis=1)

#  数据预处理
x = data.drop(['Attrition'], axis=1)
y = data['Attrition']
x = pd.get_dummies(x)
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)
# 标准化
scaler = StandardScaler()
x_train = scaler.fit_transform(x_train)
x_test = scaler.transform(x_test)
# # 交叉网格
# # 超参数网格
# param_grid = {
#     'n_estimators': [50, 60, 70, 75],       # 树的数量
#     'max_depth': [10, 15, 12],       # 树的最大深度
#     'min_samples_split': [10,  15,  20],       # 分裂内部节点所需的最小样本数
#     'min_samples_leaf': [1],         # 叶子节点的最小样本数
#     'bootstrap': [True, False]             # 是否有放回抽样
# }
#
# # 创建基础模型
# rf = RandomForestClassifier(random_state=42)
#
# # 创建 GridSearchCV 对象
# grid_search = GridSearchCV(estimator=rf,
#                            param_grid=param_grid,
#                            cv=5,               # 5折交叉验证
#                            scoring='accuracy', # 使用准确率评估
#                            n_jobs=-1,          # 多线程加速
#                            verbose=1)
#
# # 拟合数据
# grid_search.fit(x_train, y_train)
#
# # 输出最佳参数和最佳得分
# print("Best Parameters:", grid_search.best_params_)
# print("Best Cross-validation Score:", grid_search.best_score_)
# print('--------------------输出最佳参数和最佳得分-------------------')
# # 使用最优模型做预测
# best_rf = grid_search.best_estimator_
# test_acc = best_rf.score(x_test, y_test)
# print("Test Accuracy with Best Model:", test_acc)
# print('===================使用最优模型做预测=============================')
# 训练模型
model = RandomForestClassifier(bootstrap=False, max_depth=10, min_samples_leaf=1, min_samples_split=10, n_estimators=50, random_state=42)
model.fit(x_train, y_train)
# 模型评估
# print(model.score(x_test, y_test))
# print('-----------------------------')
# print(model.feature_importances_)
# 生成ROC曲线
# 获取预测概率（注意是 predict_proba）
data = pd.read_csv('../data/test2.csv')
data = data.drop(['StandardHours'], axis=1)
# 删除无用列
#  数据预处理
x = data.drop(['Attrition'], axis=1)
y = data['Attrition']
x = pd.get_dummies(x)

# 标准化
scaler = StandardScaler()
x = scaler.fit_transform(x)
y_pred = model.predict(x)
y_pred_prob = model.predict_proba(x)[:, 1]
fpr, tpr, thresholds = roc_curve(y, y_pred_prob)
auc = roc_auc_score(y, y_pred_prob)
plt.plot(fpr, tpr, label='ROC curve (area = %.2f)' % auc)
plt.plot([0, 1], [0, 1], 'k--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver operating characteristic example')
plt.legend(loc="lower right")
plt.show()
print(auc)
